import pandas as pd
import xgboost
import shap
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from sklearn.datasets import load_diabetes
# 加载数据集
diabetes = load_diabetes()
X = pd.DataFrame(diabetes.data, columns=diabetes.feature_names)
y = diabetes.target
# 训练 XGBoost 模型
model = xgboost.XGBRegressor()
model.fit(X, y)
# 创建 SHAP 解释器
explainer = shap.Explainer(model)
shap_values = explainer(X)
# 将 SHAP 值转换为 DataFrame
shap_df = pd.DataFrame(shap_values.values, columns=X.columns)
# 将 SHAP 值和特征名称整合为长格式数据
shap_long = shap_df.stack().reset_index()
shap_long.columns = ['样本', '特征', 'SHAP值']
# 添加特征值列（重复特征值以匹配 SHAP 值）
shap_long['特征值'] = X.values.flatten()
# 设置中文字体和样式
plt.rcParams['font.sans-serif'] = ['SimHei'] # 使用黑体
plt.rcParams['axes.unicode_minus'] = False # 解决负号显示问题
# 创建自定义颜色映射
custom_cmap = sns.color_palette("gist_yarg", as_cmap=True)
# 绘制 Beeswarm 图
plt.figure(figsize=(10, 6))
ax = sns.swarmplot( # 获取当前 Axes 对象
    x='SHAP值',
    y='特征',
    data=shap_long,
    size=3, # 点的大小
    hue='特征值', # 颜色基于特征值
    palette=custom_cmap, # 使用自定义颜色映射
    legend=False, # 不显示图例
)
# 手动调整颜色映射的阈值
vmin = shap_long['特征值'].min()
vmax = shap_long['特征值'].max()
norm = plt.Normalize(vmin=vmin, vmax=vmax)
sm = plt.cm.ScalarMappable(cmap=custom_cmap, norm=norm)
sm.set_array([])
# 添加自定义颜色条
cbar = plt.colorbar(sm, ax=ax) # 指定 ax 参数
cbar.set_label('特征值', rotation=270, labelpad=20)
# 设置标题和标签
plt.title('SHAP Beeswarm 图（按特征值渐变）', fontsize=14)
plt.xlabel('SHAP 值', fontsize=12)
plt.ylabel('特征', fontsize=12)
# 调整布局并显示图形
plt.tight_layout()
plt.show()